import { Quat, Vec2, Vec3, warn } from 'cc';
import { CompressionSpec } from './CompressionSpec';
import { BuiltInSensorType, IBuiltInSensor } from './IBuiltInSensor';
import { ISensor, ObservationType } from './ISensor';
import { ObservationSpec } from './ObservationSpec';
import { ObservationWriter } from './ObservationWriter';

/**
 * @en
 * A sensor implementation for vector observations.
 * 
 * @zh
 * 用于向量观察的传感器实现。
 */
export class VectorSensor implements ISensor, IBuiltInSensor {
	private _observations: number[] = [];
	public get observations(): number[] {
		return this._observations;
	}

	private _observationSpec: ObservationSpec = null!;
	public get observationSpec(): ObservationSpec {
		return this._observationSpec;
	}
	public set observationSpec(value: ObservationSpec) {
		this._observationSpec = value;
	}

	_name: string = null!;

	constructor(observationSize: number, name?: string, observationType = ObservationType.Default) {
		if (!name) {
			name = `VectorSensor_size${observationSize}`;
			if (observationType != ObservationType.Default) {
				name += `_${observationType}`;
			}
		}

		this._observations = [];
		this._name = name;
		this._observationSpec = ObservationSpec.vector(observationSize, observationType);
	}

	getBuiltInSensorType(): BuiltInSensorType {
		return BuiltInSensorType.VectorSensor;
	}
	getObservationSpec(): ObservationSpec {
		return this._observationSpec;
	}

	getObservations(): number[] {
		return this._observations;
	}

	write(writer: ObservationWriter): number {
		const exceptedObservations: number = this._observationSpec.shape[0];
		if (this._observations.length > exceptedObservations) {
			warn(
				`more observations ${this._observations.length} made than vector observation size (${exceptedObservations}). The observations will be truncated.`
			);
			this._observations.splice(exceptedObservations, this._observations.length - exceptedObservations);
		} else if (this._observations.length < exceptedObservations) {
			warn(
				`Fewer observations ${this._observations.length} made than vector observation size (${exceptedObservations}).The observations will be padded.`
			);

			for (let i = 0; i < exceptedObservations; i++) {
				this._observations.push(0);
			}
		}
		writer.addList(this._observations);
		return exceptedObservations;
	}
	getCompressedObservation(): Uint8Array {
		return new Uint8Array();
	}
	update(): void {
		this.clear();
	}
	reset(): void {
		this.clear();
	}
	getCompressionSpec(): CompressionSpec {
		return CompressionSpec.default();
	}
	getName(): string {
		return this._name;
	}

	private clear() {
		this._observations = [];
	}

	private addFloatObs(obs: number) {
		this._observations.push(obs);
	}

	addObservation(observation: number | Vec3 | Vec2 | number[] | Quat | boolean) {
		if (typeof observation === 'number') {
			this.addFloatObs(observation);
		} else if (typeof observation === 'boolean') {
			this.addFloatObs(observation ? 1 : 0);
		} else if (Array.isArray(observation)) {
			observation.forEach((e) => {
				this.addFloatObs(e);
			});
		} else if (observation instanceof Vec3) {
			this.addFloatObs(observation.x);
			this.addFloatObs(observation.y);
			this.addFloatObs(observation.z);
		} else if (observation instanceof Vec2) {
			this.addFloatObs(observation.x);
			this.addFloatObs(observation.y);
		} else if (observation instanceof Quat) {
			this.addFloatObs(observation.x);
			this.addFloatObs(observation.y);
			this.addFloatObs(observation.z);
			this.addFloatObs(observation.w);
		}
	}

	/**
	 * 添加离散分类中的独热编码( one-hot encoding) 观察室
	 * 分类标签二进制向量化，距离相等。https://zhuanlan.zhihu.com/p/634296763
	 * @param observation 观测的下标 
	 * @param range 观测的上限
	 * 比如人、 猫， 狗 共3类 range为3
	 * 人下标 observation为0 结果为 [1,0,0]的向量
	 * 猫的下标 observation为1 结果为 [0,1,0]
	 * 狗的下标 observation为2 结果为 [0,0,1]
	 */
	addOneHotObservation(observation: number, range: number) {
		for (let i = 0; i < range; i++) {
			this.addFloatObs(i == observation ? 1 : 0);
		}
	}
}
